// src/formatter.rs

use crate::parser::{format_error, Rule, SysYParser};
use pest::iterators::Pair;
use pest::Parser; // FIXED: Re-added the necessary import for SysYParser::parse

pub struct Formatter {
    output: String,
    indent_level: usize,
    is_start_of_line: bool,
}

impl Formatter {
    pub fn new() -> Self {
        Formatter {
            output: String::new(),
            indent_level: 0,
            is_start_of_line: true,
        }
    }

    pub fn format(input: &str) -> Result<String, String> {
        match SysYParser::parse(Rule::program, input) {
            Ok(pairs) => {
                let mut formatter = Formatter::new();
                formatter.format_pair(pairs.into_iter().next().unwrap());
                Ok(formatter.output.trim_end().to_string() + "\n")
            }
            Err(e) => Err(format_error(e)),
        }
    }
}

impl Formatter {
    fn write(&mut self, s: &str) {
        if self.is_start_of_line {
            self.output.push_str(&"    ".repeat(self.indent_level));
            self.is_start_of_line = false;
        }
        self.output.push_str(s);
    }

    fn writeln(&mut self) {
        self.output = self.output.trim_end().to_string();
        self.output.push('\n');
        self.is_start_of_line = true;
    }

    fn write_space(&mut self) {
        if !self.is_start_of_line {
            self.write(" ");
        }
    }
    fn indent(&mut self) {
        self.indent_level += 1;
    }
    fn dedent(&mut self) {
        self.indent_level -= 1;
    }
}

impl Formatter {
    fn format_children(&mut self, pair: Pair<Rule>) {
        for p in pair.into_inner() {
            self.format_pair(p);
        }
    }

    fn format_pair(&mut self, pair: Pair<Rule>) {
        match pair.as_rule() {
            Rule::program | Rule::CompUnit => self.format_comp_unit(pair),
            Rule::FuncDef => self.format_func_def(pair),
            Rule::Decl => self.format_children(pair), // FIXED: Corrected syntax from Rule.Decl to Rule::Decl
            Rule::ConstDecl | Rule::VarDecl => self.format_decl(pair),
            Rule::ConstDef | Rule::VarDef => self.format_def(pair),
            Rule::Block => self.format_block(pair),
            Rule::BlockItem | Rule::Stmt => self.format_children(pair),
            Rule::AssignStmt => self.format_assign_stmt(pair),
            Rule::ExpStmt => self.format_exp_stmt(pair),
            Rule::IfStmt => self.format_if_stmt(pair),
            Rule::WhileStmt => self.format_while_stmt(pair),
            Rule::ReturnStmt => self.format_return_stmt(pair),
            Rule::BreakStmt => {
                self.write("break;");
                self.writeln();
            }
            Rule::ContinueStmt => {
                self.write("continue;");
                self.writeln();
            }
            Rule::AddExp | Rule::MulExp | Rule::RelExp | Rule::EqExp | Rule::LAndExp | Rule::LOrExp => self.format_binary_exp(pair),
            Rule::UnaryExp => self.format_unary_exp(pair),
            Rule::PrimaryExp => self.format_primary_exp(pair),
            Rule::LVal => self.format_lval(pair),
            Rule::InitVal | Rule::ConstInitVal => self.format_init_val(pair),
            Rule::FuncFParams | Rule::FuncRParams => self.format_param_list(pair),
            Rule::Cond | Rule::Exp | Rule::ConstExp | Rule::FuncFParam | Rule::BType | Rule::FuncType | Rule::AddOp | Rule::MulOp | Rule::RelOp | Rule::EqOp | Rule::UnaryOp => self.format_children(pair),
            Rule::Ident | Rule::Number | Rule::INT | Rule::VOID | Rule::CONST | Rule::IF | Rule::ELSE | Rule::WHILE | Rule::BREAK | Rule::CONTINUE | Rule::RETURN | Rule::PLUS | Rule::MINUS | Rule::NOT | Rule::MUL | Rule::DIV | Rule::MOD | Rule::ASSIGN | Rule::EQ | Rule::NEQ | Rule::LT | Rule::GT | Rule::LE | Rule::GE | Rule::AND | Rule::OR => self.write(pair.as_str()),
            _ => (),
        }
    }

    fn format_comp_unit(&mut self, pair: Pair<Rule>) {
        let mut inner = pair.into_inner().peekable();
        while let Some(p) = inner.next() {
            let rule = p.as_rule();
            self.format_pair(p);
            if (rule == Rule::FuncDef || rule == Rule::Decl) && inner.peek().is_some() {
                self.writeln();
            }
        }
    }

    fn format_func_def(&mut self, pair: Pair<Rule>) {
        self.format_pair(pair.clone().into_inner().find(|p| p.as_rule() == Rule::FuncType).unwrap());
        self.write_space();
        self.format_pair(pair.clone().into_inner().find(|p| p.as_rule() == Rule::Ident).unwrap());
        self.write("(");
        if let Some(params) = pair.clone().into_inner().find(|p| p.as_rule() == Rule::FuncFParams) {
            self.format_pair(params);
        }
        self.write(")");
        self.write_space();
        self.format_pair(pair.clone().into_inner().find(|p| p.as_rule() == Rule::Block).unwrap());
        self.writeln();
    }

    fn format_decl(&mut self, pair: Pair<Rule>) {
        let mut inner = pair.into_inner().peekable();
        while let Some(p) = inner.next() {
            if p.as_rule() == Rule::COMMA {
                self.write(", ");
            } else if p.as_rule() != Rule::SEMICOLON {
                self.format_pair(p);
                if inner.peek().is_some() && inner.peek().unwrap().as_rule() != Rule::SEMICOLON {
                    self.write_space();
                }
            }
        }
        self.write(";");
        self.writeln();
    }
    
    fn format_def(&mut self, pair: Pair<Rule>) {
        let mut inner = pair.into_inner();
        self.format_pair(inner.next().unwrap()); // Ident
        while inner.peek().is_some() && inner.peek().unwrap().as_rule() == Rule::LBRACK {
            self.write("[");
            inner.next(); 
            self.format_pair(inner.next().unwrap());
            self.write("]");
            inner.next();
        }
        if inner.next().is_some() { // This consumes the ASSIGN
            self.write(" = ");
            self.format_pair(inner.next().unwrap());
        }
    }

    fn format_block(&mut self, pair: Pair<Rule>) {
        if pair.clone().into_inner().peek().is_none() {
            self.write("{}");
            return;
        }
        self.write("{");
        self.writeln();
        self.indent();
        self.format_children(pair);
        self.dedent();
        self.write("}");
    }

    fn format_assign_stmt(&mut self, pair: Pair<Rule>) {
        let mut inner = pair.into_inner();
        self.format_pair(inner.next().unwrap()); // LVal
        self.write(" = ");
        self.format_pair(inner.next().unwrap()); // Exp
        self.write(";");
        self.writeln();
    }

    fn format_exp_stmt(&mut self, pair: Pair<Rule>) {
        if let Some(exp) = pair.into_inner().find(|p| p.as_rule() == Rule::Exp) {
            self.format_pair(exp);
        }
        self.write(";");
        self.writeln();
    }

    fn format_if_stmt(&mut self, pair: Pair<Rule>) {
        self.write("if (");
        self.format_pair(pair.clone().into_inner().find(|p| p.as_rule() == Rule::Cond).unwrap());
        self.write(") ");

        let mut stmts = pair.into_inner().filter(|p| p.as_rule() == Rule::Stmt);
        let then_stmt = stmts.next().unwrap();
        let then_is_block = then_stmt.clone().into_inner().next().unwrap().as_rule() == Rule::Block;

        self.format_pair(then_stmt.clone());

        if let Some(else_stmt) = stmts.next() {
            if then_is_block {
                self.write_space();
            } else {
                self.writeln();
            }
            self.write("else ");
            self.format_pair(else_stmt);
        }
    }
    
    fn format_while_stmt(&mut self, pair: Pair<Rule>) {
        self.write("while (");
        self.format_pair(pair.clone().into_inner().find(|p| p.as_rule() == Rule::Cond).unwrap());
        self.write(") ");
        let stmt_pair = pair.into_inner().find(|p| p.as_rule() == Rule::Stmt).unwrap();
        let is_block = stmt_pair.clone().into_inner().next().unwrap().as_rule() == Rule::Block;
        if !is_block {
            self.writeln();
            self.indent();
        }
        self.format_pair(stmt_pair);
        if !is_block {
            self.dedent();
        }
    }

    fn format_return_stmt(&mut self, pair: Pair<Rule>) {
        self.write("return");
        if let Some(exp) = pair.into_inner().find(|p| p.as_rule() == Rule::Exp) {
            self.write_space();
            self.format_pair(exp);
        }
        self.write(";");
        self.writeln();
    }

    fn format_binary_exp(&mut self, pair: Pair<Rule>) {
        let mut inner = pair.into_inner();
        self.format_pair(inner.next().unwrap());
        while let Some(op) = inner.next() {
            self.write_space();
            self.format_pair(op);
            self.write_space();
            self.format_pair(inner.next().unwrap());
        }
    }

    fn format_unary_exp(&mut self, pair: Pair<Rule>) {
        // Your previous logic was better, let's restore it and make it safe
        let mut children = pair.clone().into_inner();
        let first = children.next().unwrap();
        if first.as_rule() == Rule::Ident {
            // Check if it's a function call by looking for parenthesis
            if let Some(second) = children.next() {
                if second.as_rule() == Rule::LPAREN {
                    self.format_pair(first);
                    self.write("(");
                    if let Some(params) = pair.into_inner().find(|p| p.as_rule() == Rule::FuncRParams) {
                        self.format_pair(params);
                    }
                    self.write(")");
                    return; // Done with this rule
                }
            }
        }
        // If not a function call, just format all children
        self.format_children(pair);
    }

    fn format_primary_exp(&mut self, pair: Pair<Rule>) {
        let inner = pair.into_inner().next().unwrap();
        if inner.as_rule() == Rule::Exp {
            self.write("(");
            self.format_pair(inner);
            self.write(")");
        } else {
            self.format_pair(inner);
        }
    }

    fn format_lval(&mut self, pair: Pair<Rule>) {
        let mut inner = pair.into_inner();
        self.format_pair(inner.next().unwrap()); // Ident
        for exp in inner {
            self.write("[");
            self.format_pair(exp);
            self.write("]");
        }
    }

    // FIXED: Corrected the ownership (move) error
    fn format_init_val(&mut self, pair: Pair<Rule>) {
        if pair.as_str().starts_with('{') {
            self.write("{ ");
            // Check if there are any children inside the braces before calling another function
            if pair.clone().into_inner().peekable().peek().is_some() {
                 self.format_param_list(pair);
            }
            self.write(" }");
        } else {
            self.format_children(pair);
        }
    }

    fn format_param_list(&mut self, pair: Pair<Rule>) {
        let mut inner = pair.into_inner().peekable();
        while let Some(p) = inner.next() {
            if p.as_rule() != Rule::COMMA && p.as_rule() != Rule::LBRACE && p.as_rule() != Rule::RBRACE {
                self.format_pair(p);
                if inner.peek().is_some() && inner.peek().unwrap().as_rule() != Rule::RBRACE {
                    self.write(", ");
                }
            }
        }
    }
}